#!/usr/bin/env python

# 基于pygame的键盘控制器，解决pynput在MuJoCo窗口中无法捕获键盘事件的问题

import pygame
from gym_hil.wrappers.intervention_utils import InputController

class PygameKeyboardController(InputController):
    """使用pygame捕获键盘事件的控制器，适用于MuJoCo窗口环境"""

    def __init__(self, x_step_size=0.01, y_step_size=0.01, z_step_size=0.01):
        """
        初始化控制器
        
        Args:
            x_step_size: X轴移动步长（米）
            y_step_size: Y轴移动步长（米）
            z_step_size: Z轴移动步长（米）
        """
        super().__init__(x_step_size, y_step_size, z_step_size)
        self.key_states = {
            "forward_x": False,  # 上箭头
            "backward_x": False,  # 下箭头
            "forward_y": False,  # 左箭头
            "backward_y": False,  # 右箭头
            "forward_z": False,  # 右Shift
            "backward_z": False,  # 左Shift
            "success": False,  # Enter
            "failure": False,  # Backspace
            "intervention": False,  # 空格
            "rerecord": False,  # R键
        }
        self.pygame_initialized = False
        
    def start(self):
        """初始化pygame并启动键盘监听"""
        # 初始化pygame，仅初始化事件子系统，不创建窗口
        if not pygame.get_init():
            pygame.init()
            self.pygame_initialized = True
        
        # 不创建窗口，只初始化事件子系统
        pygame.event.set_grab(False)  # 不捕获鼠标
        pygame.event.set_blocked(None)  # 不阻止任何事件
        pygame.event.set_allowed([pygame.KEYDOWN, pygame.KEYUP])  # 只允许键盘事件
            
        print("Pygame Keyboard controls:")
        print("  Arrow keys: Move in X-Y plane")
        print("  Shift and Right Shift: Move in Z axis")
        print("  Right Ctrl and Left Ctrl: Open and close gripper")
        print("  Enter: End episode with SUCCESS")
        print("  Backspace: End episode with FAILURE")
        print("  Space: Start/Stop Intervention")
        print("  ESC: Exit")
        
    def stop(self):
        """停止pygame"""
        if self.pygame_initialized:
            pygame.quit()
            self.pygame_initialized = False
            
    def update(self):
        """处理pygame事件并更新按键状态"""
        if not pygame.get_init():
            return
            
        # 处理所有待处理的事件
        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                self._handle_key_down(event.key)
            elif event.type == pygame.KEYUP:
                self._handle_key_up(event.key)
                
    def _handle_key_down(self, key):
        """处理按键按下事件"""
        if key == pygame.K_UP:
            self.key_states["forward_x"] = True
        elif key == pygame.K_DOWN:
            self.key_states["backward_x"] = True
        elif key == pygame.K_LEFT:
            self.key_states["forward_y"] = True
        elif key == pygame.K_RIGHT:
            self.key_states["backward_y"] = True
        elif key == pygame.K_LSHIFT:
            self.key_states["backward_z"] = True
        elif key == pygame.K_RSHIFT:
            self.key_states["forward_z"] = True
        elif key == pygame.K_RCTRL:
            self.open_gripper_command = True
        elif key == pygame.K_LCTRL:
            self.close_gripper_command = True
        elif key == pygame.K_RETURN:
            self.key_states["success"] = True
            self.episode_end_status = "success"
        elif key == pygame.K_BACKSPACE:
            self.key_states["failure"] = True
            self.episode_end_status = "failure"
        elif key == pygame.K_SPACE:
            self.key_states["intervention"] = not self.key_states["intervention"]
            self.intervention_flag = self.key_states["intervention"]
        elif key == pygame.K_r:
            self.key_states["rerecord"] = True
        elif key == pygame.K_ESCAPE:
            # 退出程序
            self.episode_end_status = "failure"
            
    def _handle_key_up(self, key):
        """处理按键释放事件"""
        if key == pygame.K_UP:
            self.key_states["forward_x"] = False
        elif key == pygame.K_DOWN:
            self.key_states["backward_x"] = False
        elif key == pygame.K_LEFT:
            self.key_states["forward_y"] = False
        elif key == pygame.K_RIGHT:
            self.key_states["backward_y"] = False
        elif key == pygame.K_LSHIFT:
            self.key_states["backward_z"] = False
        elif key == pygame.K_RSHIFT:
            self.key_states["forward_z"] = False
        elif key == pygame.K_RCTRL:
            self.open_gripper_command = False
        elif key == pygame.K_LCTRL:
            self.close_gripper_command = False
            
    def get_deltas(self):
        """获取当前移动增量"""
        # 计算X轴移动
        delta_x = 0.0
        if self.key_states["forward_x"]:
            delta_x += self.x_step_size
        if self.key_states["backward_x"]:
            delta_x -= self.x_step_size
            
        # 计算Y轴移动
        delta_y = 0.0
        if self.key_states["forward_y"]:
            delta_y += self.y_step_size
        if self.key_states["backward_y"]:
            delta_y -= self.y_step_size
            
        # 计算Z轴移动
        delta_z = 0.0
        if self.key_states["forward_z"]:
            delta_z += self.z_step_size
        if self.key_states["backward_z"]:
            delta_z -= self.z_step_size
            
        return delta_x, delta_y, delta_z
        
    def reset(self):
        """重置控制器状态"""
        for key in self.key_states:
            self.key_states[key] = False
        self.episode_end_status = None
        self.intervention_flag = False
        self.open_gripper_command = False
        self.close_gripper_command = False
